// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <melon/range.h>
#include <melon/io/io_buf.h>

#include <array>
#include <string>

#include <pollux/common/base/gtest_macros.h>
#include <pollux/common/base/status.h>

namespace kumo::pollux::encoding {
    class Base64 {
    public:
        static const size_t kCharsetSize = 64;
        static const size_t kReverseIndexSize = 256;

        /// Character set used for Base64 encoding.
        /// Contains specific characters that form the encoding scheme.
        using Charset = std::array<char, kCharsetSize>;

        /// Reverse lookup table for decoding.
        /// Maps each possible encoded character to its corresponding numeric value
        /// within the encoding base.
        using ReverseIndex = std::array<uint8_t, kReverseIndexSize>;

        /// Encodes the specified number of characters from the 'input'.
        static std::string encode(const char *input, size_t inputSize);

        /// Encodes the specified text.
        static std::string encode(melon::StringPiece text);

        /// Encodes the specified IOBuf data.
        static std::string encode(const melon::IOBuf *inputBuffer);

        /// Encodes the specified number of characters from the 'input' and writes the
        /// result to the 'outputBuffer'. The output must have enough space as
        /// returned by the calculateEncodedSize().
        static void encode(const char *input, size_t inputSize, char *outputBuffer);

        /// Encodes the specified number of characters from the 'input' using URL
        /// encoding.
        static std::string encodeUrl(const char *input, size_t inputSize);

        /// Encodes the specified text using URL encoding.
        static std::string encodeUrl(melon::StringPiece text);

        /// Encodes the specified IOBuf data using URL encoding.
        static std::string encodeUrl(const melon::IOBuf *inputBuffer);

        /// Encodes the specified number of characters from the 'input' and writes the
        /// result to the 'outputBuffer' using URL encoding. The output must have
        /// enough space as returned by the calculateEncodedSize().
        static void
        encodeUrl(const char *input, size_t inputSize, char *outputBuffer);

        /// Decodes the input Base64 encoded string.
        static std::string decode(melon::StringPiece encodedText);

        /// Decodes the specified encoded payload and writes the result to the
        /// 'output'.
        static void decode(
            const std::pair<const char *, int32_t> &payload,
            std::string &output);

        /// Decodes the specified number of characters from the 'input' and writes the
        /// result to the 'outputBuffer'. The output must have enough space as
        /// returned by the calculateDecodedSize().
        static void decode(const char *input, size_t inputSize, char *outputBuffer);

        /// Decodes the specified number of characters from the 'input' and writes the
        /// result to the 'outputBuffer'.
        static Status decode(
            const char *input,
            size_t inputSize,
            char *outputBuffer,
            size_t outputSize);

        /// Decodes the input Base64 URL encoded string.
        static std::string decodeUrl(melon::StringPiece encodedText);

        /// Decodes the specified URL encoded payload and writes the result to the
        /// 'output'.
        static void decodeUrl(
            const std::pair<const char *, int32_t> &payload,
            std::string &output);

        /// Decodes the specified number of characters from the 'input' using URL
        /// encoding and writes the result to the 'outputBuffer'
        static Status decodeUrl(
            const char *input,
            size_t inputSize,
            char *outputBuffer,
            size_t outputSize);

        /// Calculates the encoded size based on input 'inputSize'.
        static size_t calculateEncodedSize(size_t inputSize, bool withPadding = true);

        /// Calculates the decoded size based on encoded input and adjusts the input
        /// size for padding.
        static Expected<size_t> calculateDecodedSize(
            const char *input,
            size_t &inputSize);

    private:
        // Padding character used in encoding.
        static const char kPadding = '=';

        // Checks if the input Base64 string is padded.
        static inline bool isPadded(const char *input, size_t inputSize) {
            return (inputSize > 0 && input[inputSize - 1] == kPadding);
        }

        // Counts the number of padding characters in encoded input.
        static inline size_t numPadding(const char *input, size_t inputSize) {
            size_t numPadding{0};
            while (inputSize > 0 && input[inputSize - 1] == kPadding) {
                numPadding++;
                inputSize--;
            }
            return numPadding;
        }

        // Reverse lookup helper function to get the original index of a Base64
        // character.
        static Expected<uint8_t> base64ReverseLookup(
            char encodedChar,
            const ReverseIndex &reverseIndex);

        // Encodes the specified data using the provided charset.
        template<class T>
        static std::string
        encodeImpl(const T &input, const Charset &charset, bool includePadding);

        // Encodes the specified data using the provided charset.
        template<class T>
        static void encodeImpl(
            const T &input,
            const Charset &charset,
            bool includePadding,
            char *outputBuffer);

        // Decodes the specified data using the provided reverse lookup table.
        static Expected<size_t> decodeImpl(
            const char *input,
            size_t inputSize,
            char *outputBuffer,
            size_t outputSize,
            const ReverseIndex &reverseIndex);

        POLLUX_FRIEND_TEST(Base64Test, checksPadding);
        POLLUX_FRIEND_TEST(Base64Test, countsPaddingCorrectly);
    };
} // namespace kumo::pollux::encoding
